from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
import json
import subprocess
from fastapi.middleware.cors import CORSMiddleware
import os
import time
import psutil
import sys  # 新增，用于处理标准输入

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有HTTP方法
    allow_headers=["*"]  # 允许所有请求头
)


@app.post("/execute_code")
async def execute_code(request: Request):
    try:
        request_data = await request.json()
        code = request_data.get("code")  # 从JSON数据中获取代码字段，如果不存在则默认为空字符串
        test_input = request_data.get("input")  # 获取前端传来的测试用例数据，如果不存在则默认为空字符串
        if code == None:
            return JSONResponse(status_code=400, content={"result": "", "error": "代码内容不能为空"})

        # 根据代码所属语言来选择合适的文件名后缀等，这里以Python代码为例用.py后缀
        with open("temp_code.py", "w") as f:
            f.write(code)

        # 记录开始时间
        start_time = time.time()
        # 获取初始内存使用情况
        process = psutil.Process(os.getpid())
        initial_memory = process.memory_info().rss

        # 对于Python代码，修改执行方式以便正确处理标准输入，这里假设执行的是Python代码情况
        # 将测试输入写入到临时文件中模拟标准输入，此为一种可行的方式（也可以有其他方式比如通过管道等传递输入）
        with open('input.txt', 'w') as input_file:
            input_file.write(test_input)

        result = subprocess.run(
            ["python", "temp_code.py"],
            stdin=open('input.txt', 'r'),  # 将包含测试输入的文件作为标准输入提供给代码
            capture_output=True,
            text=True,
            encoding='utf-8'  # 指定编码格式为utf-8，确保能正确处理文本输入输出
        )

        # 记录结束时间
        end_time = time.time()
        # 获取结束时内存使用情况
        final_memory = process.memory_info().rss

        execution_time = round(end_time - start_time, 2)  # 保留小数点后两位
        memory_usage = round((final_memory - initial_memory) / 1024.0, 2)  # 转换为KB并保留小数点后两位

        return {
            "message": "代码运行成功",
            "result": result.stdout,
            "error": result.stderr,
            "execution_time": execution_time,  # 返回代码运行时间（保留小数点后两位）
            "memory_usage": memory_usage  # 返回代码运行占用内存大小（单位是KB，保留小数点后两位）
        }
    except json.JSONDecodeError:
        return JSONResponse(status_code=400, content={"result": "", "error": "请求体不是合法的JSON格式"})
    except FileNotFoundError as e:
        return JSONResponse(status_code=500, content={"result": "", "error": f"找不到文件相关错误: {str(e)}"})
    except subprocess.CalledProcessError as e:
        return JSONResponse(status_code=500, content={"result": "", "error": f"代码执行过程出错: {str(e)}"})
    finally:
        if os.path.exists("temp_code.py"):
            os.remove("temp_code.py")
        if os.path.exists("input.txt"):  # 新增，删除模拟标准输入的临时文件
            os.remove("input.txt")